import os
from contextlib import closing
import pymysql
import subprocess
from config import config_instance  # 导入配置实例

# 获取当前文件所在目录（假设这个Python文件在项目的某个合理位置）
basedir = os.path.abspath(os.path.dirname(__file__))
# 拼接出项目内备份文件夹的绝对路径
backup_folder = os.path.join(basedir, "backup")

db_name = "demo01"

def backup_database(name):
    if not os.path.exists(backup_folder):
        os.makedirs(backup_folder)

    backup_file_name = f"{name}.sql"
    backup_file_path = os.path.join(backup_folder, backup_file_name)

    try:
        conn_config = config_instance.backup_conn_config
        with closing(pymysql.connect(**conn_config)) as connection:
            connection.ping(reconnect=True)

            command = ['mysqldump', '-h', conn_config['host'], '-u', conn_config['user'],
                       '--password={}'.format(conn_config['password']), conn_config['database']]
            with open(backup_file_path, 'w') as output_file:
                process = subprocess.Popen(command, stdout=output_file, stderr=subprocess.PIPE, text=True)
                _, stderr = process.communicate()
                if process.returncode != 0:
                    raise Exception(f"mysqldump failed: {stderr}")

            print(f"数据库备份成功，备份文件为{backup_file_path}")
    except Exception as e:
        print(f"备份数据库出错: {e}")


def restore_database(backup_file_path):
    try:
        conn_config = config_instance.backup_conn_config
        with closing(pymysql.connect(**conn_config)) as connection:
            connection.ping(reconnect=True)

            command = ['mysql', '-h', conn_config['host'], '-u', conn_config['user'],
                       '--password={}'.format(conn_config['password']), conn_config['database']]
            with open(backup_file_path, 'r') as input_file:
                process = subprocess.Popen(command, stdin=input_file, stderr=subprocess.PIPE, text=True)
                _, stderr = process.communicate()
                if process.returncode != 0:
                    raise Exception(f"mysql restore failed: {stderr}")

            print("数据库恢复成功")
    except Exception as e:
        print(f"恢复数据库出错: {e}")


def check_restored_data():
    try:
        conn_config = config_instance.backup_conn_config
        with closing(pymysql.connect(**conn_config)) as connection:
            with connection.cursor() as cursor:
                # 获取所有表的名称
                cursor.execute("SHOW TABLES;")
                tables = cursor.fetchall()
                for table in tables:
                    table_name = table[0]  # 提取表名
                    print(f"\nChecking data in table: {table_name}")
                    # 对每个表执行查询语句
                    cursor.execute(f"SELECT * FROM `{table_name}` LIMIT 5;")
                    results = cursor.fetchall()

                    if not results:
                        print("No data found.")
                    else:
                        for row in results:
                            print(row)
    except Exception as e:
        print(f"检查恢复的数据出错: {e}")

if __name__ == "__main__":
    # 备份功能
    # backup_database(db_name)
    # 测试恢复功能
    backup_file_path = os.path.join(backup_folder, f"{db_name}.sql")
    restore_database(backup_file_path)

    # 检查恢复后的数据
    # check_restored_data()

